#!/usr/bin/perl
# This script is a modification of one distributed with the CodingQuarry package by 
# Alison Renwick Testa. It parses the GFF3 file output by CodingQuarry, so is missing
# many GFF3 keywords, i.e. it probably will need to be modified for other GFF3 files.
# 
# Generate FASTA file of the mRNAs and proteins, and a file of locations for each
# sequence. The locations can be loaded into a TCW database that has been created 
# from the mRNA or protein sequences.
#
use strict;
use warnings;
use Bio::Seq;
use Bio::SeqIO;
use Bio::DB::Fasta;

if (@ARGV < 3) {
	print "\nUsage: extractCoding.pl <genome fasta file> <GFF3 file> <short prefix>\n\n";
	exit;
}
print ("Executing extractCoding\n");
die ("Genome file '$ARGV[0]' does not exist ") if (!(-e $ARGV[0])); 
die ("GFF3 file '$ARGV[1]' does not exist ") if (!(-e $ARGV[1])); 

$| = 1;    # Flush output
my $outfile_cds = Bio::SeqIO->new( -format => 'fasta', -file => ">$ARGV[2].cds.fasta" );
my $outfile_pep = Bio::SeqIO->new( -format => 'fasta', -file => ">$ARGV[2].pep.fasta" );
my $outfile_loc = "$ARGV[2].loc.fasta";

### First, index the genome
my $file_fasta = $ARGV[0];
my $db = Bio::DB::Fasta->new($file_fasta);
print ("Genome fasta parsed\n");

### Second, parse the GFF3
my $SC_PREFIX = "SC_"; # prefix for scaffolds
my $SEQ_PREFIX = "$ARGV[2]_"; # prefix for protein sequences

# variables for output per sequence
my @mRNA;      # CDSs for current mRNA
my $mRNA_name; # current mRNA name
my $mRNA_loc;  # current mRNA location
my $frame;
# 
my @array;		# split line 
my %root;  		# file scaffold name, new scaffold name
my %locName;	# loc, name
my %dupName;	# for duplicate locations, i.e. alternative splicing, same start:end
my $cnt=1;
my $nGene=1;
my $dupLoc=0;

open GFF, "<$ARGV[1]" or die $!;
open LOC, ">$outfile_loc" or die $!;

while ( my $line = <GFF> ) {
    chomp $line;
    @array = split( "\t", $line );
    my $type = $array[2];
    next if ( $type eq 'exon' or $type eq 'UTR' );
	 
    if ( ( $type eq 'gene' ) and ( $. > 2 ) ) { 
        outputLastSeq();
        
        # initialize the next mRNA
        setName();
        @mRNA = (); # Empty the mRNA
        print "Seqids: " . $cnt . " mRNAs: " . $nGene . "\r" if ($nGene % 1000==0);
    }
    elsif ( $type eq 'gene' ) {    # First mRNA
    	setName();
    }
    elsif ( $type eq 'CDS' ) {
        my $cds_coord = $array[0] . " " . $array[3] . " " . $array[4];
        push( @mRNA, $cds_coord );
    }
}
#print the last one
outputLastSeq();

print "Seqids: " . $cnt . " mRNAs: " . $nGene . " Spliced " . $dupLoc . "\n";

close GFF; close LOC;
print "Finish extractCoding\n";

##################################################################
# if the same location, add suffix ".n" for mRNA_name
sub setName {
	
	if (!exists($root{$array[0]})) {
		$root{$array[0]} = $SC_PREFIX . $cnt;
		$cnt++;
	}
    $mRNA_loc = $root{$array[0]} . ":" . $array[3] . "-" . $array[4] . "(" . $array[6] . ")";
    
    if (!exists($locName{$mRNA_loc})) {
    	$mRNA_name = sprintf("%s%05d", $SEQ_PREFIX, $nGene);
    	$locName{$mRNA_loc} = $mRNA_name;
    	$nGene++;
    }
    else {
    	my $n=1;
    	$mRNA_name = $locName{$mRNA_loc} . "." . $n;
    	while (exists($dupName{$mRNA_name})) {
    		$n++;
    		$mRNA_name = $locName{$mRNA_loc} . "." . $n;
    	}
    	$dupName{$mRNA_name}++;
    	$dupLoc++;
    }
	$frame=$array[6];
}

##################################################################
sub outputLastSeq {
	# Collect CDSs and extract sequence of the previous mRNA
	my $mRNA_seq;
    foreach my $coord (@mRNA) {
        my @cds_coord = split( " ", $coord );
        my $cds_seq = $db->seq( $cds_coord[0], $cds_coord[1], $cds_coord[2] );
        $mRNA_seq .= $cds_seq;
    }
    my $output_nucleotide = Bio::Seq->new(
            -seq        => $mRNA_seq,
            -id         => $mRNA_name,
            -display_id => $mRNA_name,
            -alphabet   => 'dna',
	);
    if ($frame eq '-') {
        $output_nucleotide = $output_nucleotide->revcom();
    }
    my $output_protein = $output_nucleotide->translate();
    $outfile_cds->write_seq($output_nucleotide);
    $outfile_pep->write_seq($output_protein);
    print LOC "$mRNA_name $mRNA_loc\n";
}
